"""
API限流和安全审计服务
提供请求频率限制、安全审计日志和威胁检测功能
"""
import time
import json
import hashlib
import ipaddress
from typing import Dict, List, Optional, Any, Tuple
from datetime import datetime, timedelta
from collections import defaultdict, deque
from dataclasses import dataclass
from enum import Enum
import asyncio
import aiofiles
import re

from fastapi import Request, Response, HTTPException, status
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from starlette.middleware.base import BaseHTTPMiddleware
from core.config import settings
from core.logging_config import get_logger
from services.cache_service import cache_service

# 日志记录器
logger = get_logger("security_service")


class SecurityEventType(Enum):
    """安全事件类型"""
    LOGIN_SUCCESS = "login_success"
    LOGIN_FAILURE = "login_failure"
    LOGOUT = "logout"
    TOKEN_REFRESH = "token_refresh"
    PASSWORD_CHANGE = "password_change"
    PERMISSION_DENIED = "permission_denied"
    RATE_LIMIT_EXCEEDED = "rate_limit_exceeded"
    SUSPICIOUS_REQUEST = "suspicious_request"
    API_ACCESS = "api_access"
    DATA_ACCESS = "data_access"
    FILE_UPLOAD = "file_upload"
    SYSTEM_CONFIG_CHANGE = "system_config_change"


class ThreatLevel(Enum):
    """威胁等级"""
    LOW = "low"
    MEDIUM = "medium"
    HIGH = "high"
    CRITICAL = "critical"


@dataclass
class SecurityEvent:
    """安全事件数据结构"""
    event_type: SecurityEventType
    user_id: Optional[int]
    username: Optional[str]
    ip_address: str
    user_agent: str
    endpoint: str
    method: str
    details: Dict[str, Any]
    threat_level: ThreatLevel
    timestamp: datetime
    session_id: Optional[str] = None


@dataclass
class RateLimitInfo:
    """限流信息"""
    key: str
    count: int
    limit: int
    window_seconds: int
    reset_time: datetime


def serialize_rate_limit_info(rate_info: RateLimitInfo) -> dict:
    """序列化限流信息，处理datetime对象"""
    return {
        "key": rate_info.key,
        "count": rate_info.count,
        "limit": rate_info.limit,
        "window_seconds": rate_info.window_seconds,
        "reset_time": rate_info.reset_time.isoformat() if rate_info.reset_time else None
    }


class RateLimiter:
    """API限流器"""
    
    def __init__(self):
        self.limits = {
            # 默认限制：每分钟60次请求
            "default": {"requests": 60, "window": 1},
            # 认证端点：每分钟5次尝试
            "auth": {"requests": 60, "window": 1},
            # 文件上传：每分钟10次
            "upload": {"requests": 10, "window": 60},
            # API端点：每分钟60次
            "api": {"requests": 200, "window": 1},
            # 严格限制：每小时1000次
            "strict": {"requests": 200, "window": 60}
        }
        
        # IP白名单
        self.ip_whitelist = set()
        
        # 特殊用户限制（管理员可能有更高限制）
        
        # 内存限流数据结构
        self.memory_limits = {}  # 格式: {key: {"count": int, "reset_time": datetime}}
        self.cleanup_interval = 60  # 清理间隔（秒）
        self.last_cleanup = time.time()
        self.user_limits = {}
    
    def get_limit_key(self, request: Request, user_id: Optional[int] = None) -> str:
        """获取限流键"""
        # 获取客户端IP
        ip = self._get_client_ip(request)
        
        # 如果是白名单IP，使用特殊标识
        if ip in self.ip_whitelist:
            return f"whitelist:{ip}"
        
        # 根据用户ID和IP组合生成键
        if user_id:
            return f"user:{user_id}"
        return f"ip:{ip}"
    
    def _cleanup_memory_limits(self):
        """清理过期的内存限流记录"""
        current_time = time.time()
        
        # 如果清理间隔未到，跳过
        if current_time - self.last_cleanup < self.cleanup_interval:
            return
        
        self.last_cleanup = current_time
        
        # 清理过期记录
        expired_keys = []
        for key, data in self.memory_limits.items():
            if data["reset_time"] < datetime.now():
                expired_keys.append(key)
        
        for key in expired_keys:
            del self.memory_limits[key]
    
    def _get_memory_limit_count(self, key: str, endpoint: str, window_seconds: int) -> int:
        """从内存获取限流计数"""
        memory_key = f"{key}:{endpoint}"
        
        # 清理过期记录
        self._cleanup_memory_limits()
        
        # 获取当前计数
        if memory_key in self.memory_limits:
            data = self.memory_limits[memory_key]
            if data["reset_time"] > datetime.now():
                return data["count"]
            else:
                # 过期了，删除记录
                del self.memory_limits[memory_key]
        
        return 0
    
    def _set_memory_limit_count(self, key: str, endpoint: str, count: int, window_seconds: int):
        """设置内存限流计数"""
        memory_key = f"{key}:{endpoint}"
        reset_time = datetime.now() + timedelta(seconds=window_seconds)
        
        self.memory_limits[memory_key] = {
            "count": count,
            "reset_time": reset_time
        }
    
    def _get_client_ip(self, request: Request) -> str:
        """获取客户端IP地址"""
        # 优先从 X-Forwarded-For 获取
        forwarded_for = request.headers.get("X-Forwarded-For")
        if forwarded_for:
            return forwarded_for.split(",")[0].strip()
        
        # 其次从 X-Real-IP 获取
        real_ip = request.headers.get("X-Real-IP")
        if real_ip:
            return real_ip.strip()
        
        # 最后使用直接连接的IP
        return request.client.host if request.client else "unknown"
        if user_id:
            return f"user:{user_id}:{ip}"
        return f"ip:{ip}"
    
    def _get_client_ip(self, request: Request) -> str:
        """获取客户端真实IP"""
        # 检查代理头部
        forwarded_for = request.headers.get("X-Forwarded-For")
        if forwarded_for:
            return forwarded_for.split(",")[0].strip()
        
        real_ip = request.headers.get("X-Real-IP")
        if real_ip:
            return real_ip
        
        return request.client.host
    
    def get_limit_config(self, endpoint: str, user_id: Optional[int] = None) -> Dict[str, int]:
        """获取端点限流配置"""
        # 检查用户特殊限制
        if user_id and user_id in self.user_limits:
            return self.user_limits[user_id]
        
        # 根据端点路径确定限制类型
        if endpoint.startswith("/api/auth"):
            return self.limits["auth"]
        elif endpoint.startswith("/api/upload"):
            return self.limits["upload"]
        elif endpoint.startswith("/api/"):
            return self.limits["api"]
        else:
            return self.limits["default"]
    
    async def check_rate_limit(
        self, 
        request: Request, 
        endpoint: str,
        user_id: Optional[int] = None
    ) -> Tuple[bool, RateLimitInfo]:
        """检查是否超过限流（使用内存缓存）"""
        limit_config = self.get_limit_config(endpoint, user_id)
        limit_key = self.get_limit_key(request, user_id)
        
        # 获取当前计数（从内存）
        current_count = self._get_memory_limit_count(limit_key, endpoint, limit_config["window"])
        
        # 检查是否超过限制
        if current_count >= limit_config["requests"]:
            # 计算重置时间
            reset_time = datetime.now() + timedelta(seconds=limit_config["window"])
            
            rate_info = RateLimitInfo(
                key=limit_key,
                count=current_count,
                limit=limit_config["requests"],
                window_seconds=limit_config["window"],
                reset_time=reset_time
            )
            
            return False, rate_info
        
        # 增加计数（存储到内存）
        self._set_memory_limit_count(limit_key, endpoint, current_count + 1, limit_config["window"])
        
        rate_info = RateLimitInfo(
            key=limit_key,
            count=current_count + 1,
            limit=limit_config["requests"],
            window_seconds=limit_config["window"],
            reset_time=datetime.now() + timedelta(seconds=limit_config["window"])
        )
        
        return True, rate_info


class SecurityAuditor:
    """安全审计器"""
    
    def __init__(self):
        self.audit_log_file = "logs/security_audit.log"
        self.threat_detection_rules = self._load_threat_rules()
        self.blocked_ips = set()
        self.suspicious_patterns = [
            r"<script[^>]*>.*?</script>",  # XSS攻击
            r"union.*select.*from",        # SQL注入
            r"drop\s+table",               # SQL删除表
            r"exec\s*\(",                  # 命令执行
            r"\.\./",                      # 路径遍历
            r"etc/passwd",                 # 系统文件访问
        ]
    
    def _load_threat_rules(self) -> Dict[str, Any]:
        """加载威胁检测规则"""
        return {
            "failed_login_threshold": 5,      # 5分钟内失败登录次数
            "failed_login_window": 300,       # 5分钟窗口
            "request_frequency_threshold": 1000, # 每分钟请求数（调整阈值，避免误报）
            "suspicious_user_agents": [
                "sqlmap", "nikto", "nmap", "masscan", "zgrab"
            ],
            "blocked_countries": [],          # 被封锁的国家代码
            "suspicious_endpoints": [
                "/admin", "/wp-admin", "/phpmyadmin"
            ]
        }
    
    async def log_security_event(self, event: SecurityEvent):
        """记录安全事件"""
        try:
            # 异步写入日志文件
            await self._write_audit_log(event)
            
            # 检查威胁
            await self._check_threat(event)
            
            # 如果是高威胁事件，记录到系统日志
            if event.threat_level in [ThreatLevel.HIGH, ThreatLevel.CRITICAL]:
                logger.warning(f"高威胁安全事件: {event.event_type.value} - {event.details}")
            
        except Exception as e:
            logger.error(f"记录安全事件失败: {e}")
    
    async def _write_audit_log(self, event: SecurityEvent):
        """异步写入审计日志"""
        log_entry = {
            "timestamp": event.timestamp.isoformat(),
            "event_type": event.event_type.value,
            "user_id": event.user_id,
            "username": event.username,
            "ip_address": event.ip_address,
            "user_agent": event.user_agent,
            "endpoint": event.endpoint,
            "method": event.method,
            "threat_level": event.threat_level.value,
            "session_id": event.session_id,
            "details": event.details
        }
        
        # 确保日志目录存在
        import os
        os.makedirs(os.path.dirname(self.audit_log_file), exist_ok=True)
        
        # 异步写入文件
        async with aiofiles.open(self.audit_log_file, mode="a", encoding="utf-8") as f:
            await f.write(json.dumps(log_entry, ensure_ascii=False) + "\n")
    
    async def _check_threat(self, event: SecurityEvent):
        """检查威胁"""
        # 检查失败登录
        if event.event_type == SecurityEventType.LOGIN_FAILURE:
            await self._check_failed_login_attempts(event)
        
        # 检查可疑请求模式
        if event.event_type == SecurityEventType.SUSPICIOUS_REQUEST:
            await self._analyze_suspicious_request(event)
        
        # 检查高频请求
        await self._check_request_frequency(event)
    
    async def _check_failed_login_attempts(self, event: SecurityEvent):
        """检查失败登录尝试"""
        cache_key = f"failed_logins:{event.ip_address}"
        
        # 获取失败次数
        failed_count = await cache_service.get(cache_key)
        if failed_count is None:
            failed_count = 0
        
        failed_count += 1
        await cache_service.set(cache_key, failed_count, ttl=300)
        
        # 如果超过阈值，封锁IP
        if failed_count >= self.threat_detection_rules["failed_login_threshold"]:
            self.blocked_ips.add(event.ip_address)
            logger.warning(f"IP {event.ip_address} 因多次失败登录被临时封锁")
            
            # 记录威胁事件
            threat_event = SecurityEvent(
                event_type=SecurityEventType.SUSPICIOUS_REQUEST,
                user_id=None,
                username=None,
                ip_address=event.ip_address,
                user_agent=event.user_agent,
                endpoint=event.endpoint,
                method=event.method,
                details={
                    "reason": "failed_login_threshold_exceeded",
                    "failed_attempts": failed_count
                },
                threat_level=ThreatLevel.HIGH,
                timestamp=datetime.now()
            )
            await self.log_security_event(threat_event)
    
    async def _analyze_suspicious_request(self, event: SecurityEvent):
        """分析可疑请求"""
        # 检查用户代理
        for suspicious_agent in self.threat_detection_rules["suspicious_user_agents"]:
            if suspicious_agent.lower() in event.user_agent.lower():
                self.blocked_ips.add(event.ip_address)
                logger.warning(f"检测到可疑用户代理: {event.user_agent}")
                break
        
        # 检查请求参数中的可疑模式
        for pattern in self.suspicious_patterns:
            if re.search(pattern, str(event.details), re.IGNORECASE):
                logger.warning(f"检测到可疑请求模式: {pattern}")
                
                # 记录高威胁事件
                threat_event = SecurityEvent(
                    event_type=SecurityEventType.SUSPICIOUS_REQUEST,
                    user_id=event.user_id,
                    username=event.username,
                    ip_address=event.ip_address,
                    user_agent=event.user_agent,
                    endpoint=event.endpoint,
                    method=event.method,
                    details={
                        "reason": "suspicious_pattern_detected",
                        "pattern": pattern
                    },
                    threat_level=ThreatLevel.HIGH,
                    timestamp=datetime.now()
                )
                await self.log_security_event(threat_event)
    
    async def _check_request_frequency(self, event: SecurityEvent):
        """检查请求频率"""
        # 排除静态资源请求，避免误报
        if event.endpoint and (
            event.endpoint.startswith('/static/') or
            event.endpoint.endswith(('.css', '.js', '.png', '.jpg', '.jpeg', '.gif', '.svg', '.ico', '.woff', '.woff2', '.ttf'))
        ):
            return  # 静态资源不参与频率限制
        
        cache_key = f"request_freq:{event.ip_address}"
        
        # 获取请求计数
        request_count = await cache_service.get(cache_key)
        if request_count is None:
            request_count = 0
        
        request_count += 1
        await cache_service.set(cache_key, request_count, ttl=60)
        
        # 检查是否超过阈值
        if request_count >= self.threat_detection_rules["request_frequency_threshold"]:
            logger.warning(f"IP {event.ip_address} 请求频率过高: {request_count}/分钟,endpoint:{event.endpoint}")
            
            # 记录威胁事件
            threat_event = SecurityEvent(
                event_type=SecurityEventType.SUSPICIOUS_REQUEST,
                user_id=event.user_id,
                username=event.username,
                ip_address=event.ip_address,
                user_agent=event.user_agent,
                endpoint=event.endpoint,
                method=event.method,
                details={
                    "reason": "high_request_frequency",
                    "requests_per_minute": request_count
                },
                threat_level=ThreatLevel.MEDIUM,
                timestamp=datetime.now()
            )
            await self.log_security_event(threat_event)
    
    def is_ip_blocked(self, ip_address: str) -> bool:
        """检查IP是否被封锁"""
        return ip_address in self.blocked_ips
    
    def analyze_request(self, request: Request) -> Optional[Dict[str, Any]]:
        """分析请求安全性"""
        analysis = {
            "is_suspicious": False,
            "threat_level": ThreatLevel.LOW,
            "reasons": []
        }
        
        # 检查IP封锁
        ip = self._get_client_ip(request)
        if self.is_ip_blocked(ip):
            analysis["is_suspicious"] = True
            analysis["threat_level"] = ThreatLevel.HIGH
            analysis["reasons"].append("ip_blocked")
            return analysis
        
        # 检查用户代理
        user_agent = request.headers.get("User-Agent", "")
        for suspicious_agent in self.threat_detection_rules["suspicious_user_agents"]:
            if suspicious_agent.lower() in user_agent.lower():
                analysis["is_suspicious"] = True
                analysis["threat_level"] = ThreatLevel.MEDIUM
                analysis["reasons"].append("suspicious_user_agent")
        
        # 检查端点
        endpoint = request.url.path
        for suspicious_endpoint in self.threat_detection_rules["suspicious_endpoints"]:
            if suspicious_endpoint in endpoint:
                analysis["is_suspicious"] = True
                analysis["threat_level"] = ThreatLevel.MEDIUM
                analysis["reasons"].append("suspicious_endpoint")
        
        return analysis if analysis["is_suspicious"] else None
    
    def _get_client_ip(self, request: Request) -> str:
        """获取客户端真实IP"""
        forwarded_for = request.headers.get("X-Forwarded-For")
        if forwarded_for:
            return forwarded_for.split(",")[0].strip()
        
        real_ip = request.headers.get("X-Real-IP")
        if real_ip:
            return real_ip
        
        return request.client.host


class SecurityMiddleware(BaseHTTPMiddleware):
    """安全中间件"""
    
    def __init__(self, app):
        super().__init__(app)
        self.rate_limiter = RateLimiter()
        self.security_auditor = SecurityAuditor()
    
    async def dispatch(self, request: Request, call_next):
        """处理请求"""
        start_time = time.time()
        
        # 获取用户信息
        user_id = None
        username = None
        
        try:
            # 从Authorization头获取用户信息
            auth_header = request.headers.get("Authorization")
            if auth_header and auth_header.startswith("Bearer "):
                # 这里可以解析JWT获取用户信息
                # 为了性能，这里简化处理
                pass
        except Exception:
            pass
        
        # 获取客户端信息
        ip_address = self._get_client_ip(request)
        user_agent = request.headers.get("User-Agent", "")
        endpoint = request.url.path
        method = request.method
        
        # 安全检查
        security_analysis = self.security_auditor.analyze_request(request)
        if security_analysis:
            # 记录可疑请求
            event = SecurityEvent(
                event_type=SecurityEventType.SUSPICIOUS_REQUEST,
                user_id=user_id,
                username=username,
                ip_address=ip_address,
                user_agent=user_agent,
                endpoint=endpoint,
                method=method,
                details=security_analysis,
                threat_level=security_analysis["threat_level"],
                timestamp=datetime.now()
            )
            await self.security_auditor.log_security_event(event)
            
            # 如果是高威胁，拒绝请求
            if security_analysis["threat_level"] == ThreatLevel.HIGH:
                raise HTTPException(
                    status_code=status.HTTP_403_FORBIDDEN,
                    detail="请求被安全系统阻止"
                )
        
        # 限流检查
        rate_limit_allowed, rate_info = await self.rate_limiter.check_rate_limit(
            request, endpoint, user_id
        )
        
        if not rate_limit_allowed:
            # 记录限流事件
            event = SecurityEvent(
                event_type=SecurityEventType.RATE_LIMIT_EXCEEDED,
                user_id=user_id,
                username=username,
                ip_address=ip_address,
                user_agent=user_agent,
                endpoint=endpoint,
                method=method,
                details={
                    "rate_limit_info": serialize_rate_limit_info(rate_info)
                },
                threat_level=ThreatLevel.MEDIUM,
                timestamp=datetime.now()
            )
            await self.security_auditor.log_security_event(event)
            logger.info(f"限流检查结果 - 允许: {rate_limit_allowed}, 限制信息: {rate_info} request[{request}] endpoint[{endpoint}] userid[{user_id}]")
            raise HTTPException(
                status_code=status.HTTP_429_TOO_MANY_REQUESTS,
                detail=f"请求频率过高，请在 {rate_info.reset_time.strftime('%H:%M:%S')} 后重试"
            )
        
        # 添加限流信息到响应头
        response = await call_next(request)
        response.headers["X-RateLimit-Limit"] = str(rate_info.limit)
        response.headers["X-RateLimit-Remaining"] = str(rate_info.limit - rate_info.count)
        response.headers["X-RateLimit-Reset"] = rate_info.reset_time.isoformat()
        
        # 记录API访问事件
        if endpoint.startswith("/api/"):
            event = SecurityEvent(
                event_type=SecurityEventType.API_ACCESS,
                user_id=user_id,
                username=username,
                ip_address=ip_address,
                user_agent=user_agent,
                endpoint=endpoint,
                method=method,
                details={
                    "response_time": time.time() - start_time,
                    "status_code": response.status_code
                },
                threat_level=ThreatLevel.LOW,
                timestamp=datetime.now()
            )
            await self.security_auditor.log_security_event(event)
        
        return response
    
    def _get_client_ip(self, request: Request) -> str:
        """获取客户端真实IP"""
        forwarded_for = request.headers.get("X-Forwarded-For")
        if forwarded_for:
            return forwarded_for.split(",")[0].strip()
        
        real_ip = request.headers.get("X-Real-IP")
        if real_ip:
            return real_ip
        
        return request.client.host


# 全局实例
rate_limiter = RateLimiter()
security_auditor = SecurityAuditor()


async def require_rate_limit(request: Request, endpoint: str = None, user_id: int = None):
    """限流装饰器"""
    if endpoint is None:
        endpoint = request.url.path
    
    allowed, rate_info = await rate_limiter.check_rate_limit(request, endpoint, user_id)
    logger.debug(f"限流检查结果 - 允许: {allowed}, 重置时间: {rate_info.reset_time} request[{request}] endpoint[{endpoint}] userid[{user_id}]")
    if not allowed:
        raise HTTPException(
            status_code=status.HTTP_429_TOO_MANY_REQUESTS,
            detail=f"请求频率过高，请在 {rate_info.reset_time.strftime('%H:%M:%S')} 后重试"
        )
    
    return rate_info


async def log_security_event(
    event_type: SecurityEventType,
    request: Request,
    details: Dict[str, Any] = None,
    user_id: int = None,
    username: str = None,
    threat_level: ThreatLevel = ThreatLevel.LOW
):
    """记录安全事件"""
    event = SecurityEvent(
        event_type=event_type,
        user_id=user_id,
        username=username,
        ip_address=security_auditor._get_client_ip(request),
        user_agent=request.headers.get("User-Agent", ""),
        endpoint=request.url.path,
        method=request.method,
        details=details or {},
        threat_level=threat_level,
        timestamp=datetime.now()
    )
    
    await security_auditor.log_security_event(event)


async def get_security_statistics() -> Dict[str, Any]:
    """获取安全统计信息"""
    try:
        # 获取最近24小时的安全事件
        now = datetime.now()
        yesterday = now - timedelta(days=1)
        
        # 这里可以从缓存或数据库获取统计数据
        # 为了演示，返回示例数据
        return {
            "total_events": 0,
            "events_by_type": {},
            "events_by_threat_level": {},
            "blocked_ips": len(security_auditor.blocked_ips),
            "rate_limit_violations": 0,
            "failed_login_attempts": 0,
            "suspicious_requests": 0,
            "period": "24h"
        }
    except Exception as e:
        logger.error(f"获取安全统计信息失败: {e}")
        return {}